import queue

import rospy
import yaml
import IDB.imudb
import torch
import threading
import time
import numpy as np
from copy import deepcopy
from threading import Thread
from collections import deque
from sensor_msgs.msg import Imu


class IMU_Denoiser(Thread):

    def __init__(self, imuTopic,imuPub,config_fp, ckpts_fp, imuRate):
        super(IMU_Denoiser, self).__init__()
        self.processIMU = None
        self.publishIMU = None
        self.imuTemplate = []
        self.processIMU=None
        self.publishIMU=None
        self.imuBuffer=queue.Queue()
        self.IMU_Processed = queue.Queue()
        self.imusub = rospy.Subscriber(imuTopic, Imu, self.handleIMU, queue_size=200)
        self.imupub = rospy.Publisher(imuPub, Imu, queue_size=500)
        self.config = None
        self.model = None
        self.backend = None
        self.x_imu_buffer = None

        # self. REJECT_MSE_THRE = 0.001
        # DIY
        self. REJECT_MSE_THRE = 0.001
        self.ROUND_DIGIT = 6
        self.reject_cnt = 0
        self.msg_cnt = 0
        self.reject_cnt = 0
        self.max_loss = 0
        self.config_fp="params/log/hparams.yaml"
        self.ckpts_fp="checkpoint/IMUDB.ckpt"
        self.initParameter()
        self.lastUpdate = 0
        self.imurate = rospy.Rate(200)

    def run(self):
        self.processIMU = threading.Thread(target=self.producer)
        self.publishIMU = threading.Thread(target=self.consumer)
        self.processIMU.start()
        time.sleep(160)
        # time.sleep(15)
        self.publishIMU.start()

    def consumer(self):
        print("start sending!")
        while True:
            if self.IMU_Processed.qsize() > 0:
                item = self.IMU_Processed.get()
                self.imupub.publish(item)
                self.imurate.sleep()
            else:
                if self.isAlive:
                    #print("IMU publish thread is running")
                    time.sleep(1)
                else:
                    print("IMU Producer stopped , end this thread too")
                    return


    def producer(self):
        while True:
            list = []
            Max=0
            update=True
            current = time.perf_counter()
            if self.imuBuffer.qsize()>300:
                Max=300
            else:
                if self.lastUpdate==0 or (self.lastUpdate!=0 and  current-self.lastUpdate<9):
                    print("Not enough IMU,waiting for new data")
                    time.sleep(3)
                    continue
                else:
                    print("too long for new data")
                    if self.imuBuffer.qsize()==0:
                        print("No imu received!End this thread")
                        self.isAlive=False
                        return
                    else:
                        print("send the rest imu")
                        update=False
                        Max=self.imuBuffer.qsize()
            for i in range(0, Max):
                list.append(self.imuBuffer.get())
            list.reverse()
            #print("save imu to publish queue")
            while len(list)>0:
                self.IMU_Processed.put(list.pop())

            if update:
                self.lastUpdate = time.perf_counter()
            time.sleep(2)

    def initParameter(self):
        with open(self.config_fp) as f:
            self.config = yaml.safe_load(f)
        self.config = self.config['config']['model']
        print("Loading params ...")
        print(self.config)
        print("Initializing the ckpts_fp session ...")
        self.model = IDB.imudb.Model.load_from_checkpoint(self.ckpts_fp)
        self.backend = self.model.limu_bert_mlm.forward
        self.x_imu_buffer = deque(maxlen=int(self.config['inputs']['imu_sequence_number']))
        print("Initialize Done")

    def handleIMU(self,msg):
        self.msg_cnt=self.msg_cnt+1
        self.x_imu_buffer.append(
            [
                msg.angular_velocity.x, msg.angular_velocity.y, msg.angular_velocity.z,
                msg.linear_acceleration.x, msg.linear_acceleration.y, msg.linear_acceleration.z
            ]
        )  # (S, 6)
        new_imu_msg = deepcopy(msg)
        buffer_is_valid = (len(self.x_imu_buffer) == int(self.config['inputs']['imu_sequence_number']))

        if buffer_is_valid:
            x_imu_infer = np.array([self.x_imu_buffer]).astype(np.float32)
            x_imu_infer[:, :, 3:] = x_imu_infer[:, :, 3:] / 9.8
            hat_imu = self.backend(torch.tensor(x_imu_infer))

            hat_imu = hat_imu.detach().numpy()
            # denorm
            hat_imu = hat_imu[0]
            hat_imu[:, 3:] = hat_imu[:, 3:] * 9.8
            hat_imu_now = hat_imu[-1, :]
            mse = (np.square(self.x_imu_buffer[-1] - hat_imu_now)).mean()
            mse = np.round(mse, self.ROUND_DIGIT)
            if mse > self.REJECT_MSE_THRE:
                print(f"mse = {mse} > {self.REJECT_MSE_THRE}, Using old msg.")
                self.imuBuffer.put(msg)
                self.reject_cnt += 1
                return
            else:
                print(f"mse = {mse} < {self.REJECT_MSE_THRE}, Replacing imu.")
            if mse > self.max_loss:
                self.max_loss = mse

            new_imu_msg.angular_velocity.x = hat_imu_now[0]
            new_imu_msg.angular_velocity.y = hat_imu_now[1]
            new_imu_msg.angular_velocity.z = hat_imu_now[2]


            new_imu_msg.linear_acceleration.x = hat_imu_now[3]
            new_imu_msg.linear_acceleration.y = hat_imu_now[4]
            new_imu_msg.linear_acceleration.z = hat_imu_now[5]
            # self.imuBuffer.put(new_imu_msg)
            #print("x :" + str(msg.linear_acceleration.x - new_imu_msg.linear_acceleration.x) + "  y:" + str(msg.linear_acceleration.y - new_imu_msg.linear_acceleration.y) + "  z:" + str(msg.linear_acceleration.z - new_imu_msg.linear_acceleration.z))
            self.imuBuffer.put(msg)

